//! 分层池化系统
//! 
//! 实现冷池、温池、热池三层架构

pub mod cold_pool;
pub mod warm_pool; 
pub mod hot_pool;

#[cfg(test)]
pub mod test_utils;

use crate::Result;
use crate::types::{PoolType, Language, SystemStats};
use crate::sandbox_manager::{Sandbox, LifecycleManager};
use crate::firecracker::snapshot::SnapshotManager;
use std::collections::HashMap;
use std::time::{SystemTime, Instant};
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{info, warn, error, debug};
use tokio::time::Duration;

/// 分层池化管理器
pub struct LayeredPoolManager {
    /// 冷池管理器
    cold_pool: Arc<cold_pool::ColdPool>,
    /// 温池管理器  
    warm_pool: Arc<warm_pool::WarmPool>,
    /// 热池管理器
    hot_pool: Arc<hot_pool::HotPool>,
    /// 系统配置
    config: Arc<crate::config::PoolConfig>,
    /// 全局统计信息
    global_stats: Arc<tokio::sync::RwLock<GlobalPoolStats>>,
}

/// 全局池化统计信息
#[derive(Debug, Clone)]
pub struct GlobalPoolStats {
    pub total_requests: usize,
    pub hot_pool_hits: usize,
    pub warm_pool_hits: usize,
    pub cold_pool_hits: usize,
    pub cache_misses: usize,
    pub avg_response_time_ms: f64,
    pub language_distribution: HashMap<Language, usize>,
    pub error_count: usize,
    pub last_updated: SystemTime,
}

impl Default for GlobalPoolStats {
    fn default() -> Self {
        Self {
            total_requests: 0,
            hot_pool_hits: 0,
            warm_pool_hits: 0,
            cold_pool_hits: 0,
            cache_misses: 0,
            avg_response_time_ms: 0.0,
            language_distribution: HashMap::new(),
            error_count: 0,
            last_updated: SystemTime::now(),
        }
    }
}

impl LayeredPoolManager {
    /// 创建分层池化管理器
    pub async fn new(
        config: Arc<crate::config::PoolConfig>,
        snapshot_manager: Arc<SnapshotManager>,
        lifecycle_manager: Arc<LifecycleManager>,
    ) -> Result<Self> {
        info!("初始化分层池化管理器");
        
        // 创建冷池
        // 提取SnapshotManager
        let snapshot_manager = Arc::try_unwrap(snapshot_manager)
            .map_err(|_| crate::Error::Internal("Failed to unwrap SnapshotManager".to_string()))?;
        
        let cold_pool = Arc::new(cold_pool::ColdPool::new(
            config.clone(),
            Arc::new(Mutex::new(snapshot_manager)),
            lifecycle_manager.clone(),
        ));
        
        // 创建温池
        let warm_pool = Arc::new(warm_pool::WarmPool::new(
            config.clone(),
            cold_pool.clone(),
            lifecycle_manager.clone(),
        ));
        
        // 创建热池
        let mut hot_pool = hot_pool::HotPool::new(config.clone());
        hot_pool.set_warm_pool(warm_pool.clone());
        hot_pool.set_lifecycle_manager(lifecycle_manager.clone());
        let hot_pool = Arc::new(hot_pool);
        
        let manager = Self {
            cold_pool,
            warm_pool,
            hot_pool,
            config,
            global_stats: Arc::new(tokio::sync::RwLock::new(GlobalPoolStats::default())),
        };
        
        // 初始化冷池快照
        manager.cold_pool.initialize_snapshots().await?;
        
        // 启动维护任务
        manager.start_maintenance_tasks().await?;
        
        info!("分层池化管理器初始化完成");
        Ok(manager)
    }
    
    /// 智能获取沙箱（按优先级：热池→温池→冷池）
    pub async fn acquire_sandbox(&self, language: &Language) -> Result<Option<Sandbox>> {
        let start_time = Instant::now();
        debug!("开始智能获取沙箱 - 语言: {:?}", language);
        
        // 更新请求统计
        {
            let mut stats = self.global_stats.write().await;
            stats.total_requests += 1;
            *stats.language_distribution.entry(language.clone()).or_insert(0) += 1;
        }
        
        let mut sandbox_result: Option<Sandbox> = None;
        let mut pool_source = PoolType::Hot;
        
        // 1. 优先从热池获取（最快 <5ms）
        match self.hot_pool.acquire(language).await {
            Ok(Some(sandbox)) => {
                info!("从热池成功获取沙箱 - 语言: {:?}", language);
                sandbox_result = Some(sandbox);
                pool_source = PoolType::Hot;
            }
            Ok(None) => {
                debug!("热池无可用沙箱，尝试温池");
                
                // 2. 从温池获取并激活（中等 <50ms）
                match self.warm_pool.acquire(language).await {
                    Ok(Some(sandbox)) => {
                        info!("从温池成功获取沙箱 - 语言: {:?}", language);
                        sandbox_result = Some(sandbox);
                        pool_source = PoolType::Warm;
                    }
                    Ok(None) => {
                        debug!("温池无可用沙箱，尝试冷池");
                        
                        // 3. 从冷池创建（较慢 <200ms）
                        match self.cold_pool.acquire(language).await {
                            Ok(Some(sandbox)) => {
                                info!("从冷池成功创建沙箱 - 语言: {:?}", language);
                                sandbox_result = Some(sandbox);
                                pool_source = PoolType::Cold;
                            }
                            Ok(None) => {
                                warn!("所有池都无法提供 {:?} 沙箱", language);
                            }
                            Err(e) => {
                                error!("从冷池创建沙箱失败: {}", e);
                                let mut stats = self.global_stats.write().await;
                                stats.error_count += 1;
                                return Err(e);
                            }
                        }
                    }
                    Err(e) => {
                        error!("从温池获取沙箱失败: {}", e);
                        let mut stats = self.global_stats.write().await;
                        stats.error_count += 1;
                        return Err(e);
                    }
                }
            }
            Err(e) => {
                error!("从热池获取沙箱失败: {}", e);
                let mut stats = self.global_stats.write().await;
                stats.error_count += 1;
                return Err(e);
            }
        }
        
        // 更新统计信息
        let duration = start_time.elapsed();
        {
            let mut stats = self.global_stats.write().await;
            
            match pool_source {
                PoolType::Hot => stats.hot_pool_hits += 1,
                PoolType::Warm => stats.warm_pool_hits += 1,
                PoolType::Cold => stats.cold_pool_hits += 1,
            }
            
            if sandbox_result.is_none() {
                stats.cache_misses += 1;
            }
            
            // 更新平均响应时间
            let new_avg = (stats.avg_response_time_ms * (stats.total_requests - 1) as f64 
                + duration.as_millis() as f64) / stats.total_requests as f64;
            stats.avg_response_time_ms = new_avg;
            stats.last_updated = SystemTime::now();
        }
        
        if let Some(ref sandbox) = sandbox_result {
            info!(
                "智能获取沙箱成功 - 语言: {:?}, 来源: {:?}, 耗时: {}ms",
                language,
                pool_source,
                duration.as_millis()
            );
        } else {
            warn!(
                "智能获取沙箱失败 - 语言: {:?}, 耗时: {}ms",
                language,
                duration.as_millis()
            );
        }
        
        Ok(sandbox_result)
    }
    
    /// 智能释放沙箱回池
    pub async fn release_sandbox(&self, sandbox: Sandbox) -> Result<()> {
        let language = sandbox.language().clone();
        debug!("开始智能释放沙箱 - 语言: {:?}", language);
        
        // 检查沙箱状态，决定释放到哪个池
        if self.should_return_to_hot_pool(&sandbox).await {
            // 释放到热池（最快复用）
            debug!("释放沙箱到热池 - 语言: {:?}", language);
            self.hot_pool.release(sandbox).await
        } else if self.should_return_to_warm_pool(&sandbox).await {
            // 降级到温池（需要重新预热）
            debug!("降级沙箱到温池 - 语言: {:?}", language);
            self.warm_pool.add_prewarmed(sandbox).await
        } else {
            // 直接销毁（不适合复用）
            info!("沙箱不适合复用，直接销毁 - 语言: {:?}", language);
            // 销毁逻辑由热池的release方法处理
            self.hot_pool.release(sandbox).await
        }
    }
    
    /// 获取系统统计
    pub async fn get_system_stats(&self) -> Result<SystemStats> {
        let mut pool_stats = HashMap::new();
        
        // 收集各池统计
        pool_stats.insert(PoolType::Cold, self.cold_pool.get_stats().await?);
        pool_stats.insert(PoolType::Warm, self.warm_pool.get_stats().await?);
        pool_stats.insert(PoolType::Hot, self.hot_pool.get_stats().await?);
        
        let total_sandboxes: usize = pool_stats.values().map(|s| s.total_capacity).sum();
        let active_sandboxes: usize = pool_stats.values().map(|s| s.used_count + s.idle_count).sum();
        
        let global_stats = self.global_stats.read().await;
        
        Ok(SystemStats {
            total_sandboxes,
            active_sandboxes,
            pool_stats,
            language_distribution: global_stats.language_distribution.clone(),
            system_load: self.calculate_system_load().await,
            memory_usage_percent: self.calculate_memory_usage().await,
            timestamp: SystemTime::now(),
        })
    }
    
    /// 获取全局池化统计
    pub async fn get_global_stats(&self) -> GlobalPoolStats {
        self.global_stats.read().await.clone()
    }
    
    /// 预热指定语言的池
    pub async fn preheat_language(&self, language: Language, hot_count: usize, warm_count: usize) -> Result<()> {
        info!("预热语言池 - 语言: {:?}, 热池: {}, 温池: {}", language, hot_count, warm_count);
        
        // 并行预热温池和热池
        let warm_preheat = self.warm_pool.preheat(language.clone(), warm_count);
        let hot_preheat = self.hot_pool.preheat(&language, hot_count);
        
        // 等待预热完成
        let (warm_result, hot_result) = tokio::join!(warm_preheat, hot_preheat);
        
        if let Err(e) = warm_result {
            error!("温池预热失败: {}", e);
        }
        
        if let Err(e) = hot_result {
            error!("热池预热失败: {}", e);
        }
        
        Ok(())
    }
    
    /// 启动维护任务
    async fn start_maintenance_tasks(&self) -> Result<()> {
        info!("启动池化系统维护任务");
        
        // 启动热池维护
        self.hot_pool.start_maintenance().await?;
        
        // 启动温池自动预热
        self.warm_pool.start_auto_preheating().await?;
        
        // 启动全局清理任务
        let manager = self.clone_for_task();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60));
            
            loop {
                interval.tick().await;
                
                if let Err(e) = manager.perform_global_cleanup().await {
                    error!("全局清理任务失败: {}", e);
                }
            }
        });
        
        // 启动统计报告任务
        let stats_manager = self.clone_for_task();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(300)); // 5分钟
            
            loop {
                interval.tick().await;
                
                stats_manager.report_system_stats().await;
            }
        });
        
        Ok(())
    }
    
    /// 执行全局清理
    async fn perform_global_cleanup(&self) -> Result<()> {
        debug!("执行全局池化系统清理");
        
        // 清理各个池的过期沙箱
        let cold_cleanup = self.cold_pool.cleanup_expired_snapshots();
        let warm_cleanup = self.warm_pool.cleanup_expired_sandboxes();
        let hot_cleanup = self.hot_pool.cleanup_invalid_sandboxes();
        
        let (cold_result, warm_result, hot_result) = tokio::join!(
            cold_cleanup,
            warm_cleanup,
            hot_cleanup
        );
        
        if let Err(e) = cold_result {
            error!("冷池清理失败: {}", e);
        }
        
        if let Err(e) = warm_result {
            error!("温池清理失败: {}", e);
        }
        
        if let Err(e) = hot_result {
            error!("热池清理失败: {}", e);
        }
        
        info!("全局池化系统清理完成");
        Ok(())
    }
    
    /// 报告系统统计信息
    async fn report_system_stats(&self) {
        match self.get_system_stats().await {
            Ok(stats) => {
                let global_stats = self.global_stats.read().await;
                
                info!(
                    "池化系统统计报告 - 总请求: {}, 热池命中: {}, 温池命中: {}, 冷池命中: {}, 缓存未命中: {}, 平均响应时间: {:.2}ms",
                    global_stats.total_requests,
                    global_stats.hot_pool_hits,
                    global_stats.warm_pool_hits,
                    global_stats.cold_pool_hits,
                    global_stats.cache_misses,
                    global_stats.avg_response_time_ms
                );
                
                info!(
                    "系统资源 - 总沙箱: {}, 活跃沙箱: {}, 系统负载: {:.2}, 内存使用: {:.1}%",
                    stats.total_sandboxes,
                    stats.active_sandboxes,
                    stats.system_load,
                    stats.memory_usage_percent
                );
            }
            Err(e) => {
                error!("获取系统统计失败: {}", e);
            }
        }
    }
    
    // 私有方法
    
    /// 检查是否应该返回到热池
    async fn should_return_to_hot_pool(&self, sandbox: &Sandbox) -> bool {
        // 检查使用次数
        sandbox.metadata.usage_count < self.config.max_reuse_count.try_into().unwrap() &&
        // 检查运行时间
        sandbox.metadata.uptime()
            .map(|uptime| uptime < Duration::from_secs(self.config.max_sandbox_lifetime_seconds))
            .unwrap_or(true) &&
        // 检查沙箱状态
        matches!(sandbox.metadata.state, crate::types::SandboxState::Running)
    }
    
    /// 检查是否应该返回到温池
    async fn should_return_to_warm_pool(&self, sandbox: &Sandbox) -> bool {
        // 如果不适合热池但还可以使用，降级到温池
        sandbox.metadata.usage_count < (self.config.max_reuse_count * 2) as u64 &&
        sandbox.metadata.uptime()
            .map(|uptime| uptime < Duration::from_secs(self.config.max_sandbox_lifetime_seconds * 2))
            .unwrap_or(true)
    }
    
    /// 计算系统负载
    async fn calculate_system_load(&self) -> f64 {
        // 简单的负载计算：基于活跃任务数和池使用率
        let global_stats = self.global_stats.read().await;
        let recent_requests = global_stats.total_requests as f64;
        
        // 基于最近的命中率计算负载
        let total_hits = global_stats.hot_pool_hits + global_stats.warm_pool_hits + global_stats.cold_pool_hits;
        let hit_rate = if global_stats.total_requests > 0 {
            total_hits as f64 / global_stats.total_requests as f64
        } else {
            1.0
        };
        
        // 负载 = (1 - 命中率) * 请求频率因子
        (1.0 - hit_rate) * (recent_requests / 1000.0).min(1.0)
    }
    
    /// 计算内存使用率
    async fn calculate_memory_usage(&self) -> f64 {
        // TODO: 实现真实的内存使用率计算
        // 这里返回一个估算值
        60.0 + (self.global_stats.read().await.total_requests as f64 / 100.0).min(30.0)
    }
    
    /// 为任务克隆管理器（共享数据）
    fn clone_for_task(&self) -> LayeredPoolManager {
        LayeredPoolManager {
            cold_pool: Arc::clone(&self.cold_pool),
            warm_pool: Arc::clone(&self.warm_pool),
            hot_pool: Arc::clone(&self.hot_pool),
            config: Arc::clone(&self.config),
            global_stats: Arc::clone(&self.global_stats),
        }
    }
}